from hep.aida import *

from java.awt import Color
from java.awt import Font
from java.lang import *
from java.util import *
from java.io import *

from jhplot.jadraw import *
from jhplot  import *
from jplot  import *

from fr.in2p3.ipnl.aidautilities import *
from fr.in2p3.ipnl.oflap.scint import *
from jscinttools..utilities.jhplotutilities import *
from jscinttools.scint.models import *

Locale.setDefault(Locale.ENGLISH)
System.getProperties().setProperty("hep.aida.BatchAnalysisFactory","hep.aida.ref.BatchAnalysisFactory")

showFigures = False
writeFigures = not showFigures

af = IAnalysisFactory.create();
fitter = af.createFitFactory().createFitter("Chi2", "jminuit");

""" Define colors and fonts"""
figType = [".png", ".eps"]
colors = [Color.blue, Color.red, Color(0, 170, 0), Color.black, Color.yellow]
axisLabelFont = Font(Font.SERIF, Font.PLAIN, 25)
titleFont = Font(Font.SERIF, Font.PLAIN, 25)
legendFont = Font(Font.SERIF, Font.PLAIN, 20)
ticFont = Font(Font.SERIF, Font.PLAIN, 20)
statBoxFont = Font(Font.SERIF, Font.PLAIN, 20)

""" Define the Path and files names"""
aidaPathL = "/home/marc-antoine/Documents/BGO/QFin"
aidaFileL = "ScintillationAnalysis_BGO"
aidaPathQ = "/home/marc-antoine/Documents/BGO/QFin/QueensData/queenstools"
aidaFileQ = "ScintillationAnalysis_BGO_QFile"
outPath = "/home/marc-antoine/Dropbox/BGO"


"""
###############################################################
# Decay Models
###############################################################
"""

modelAgLvl = 3

if modelAgLvl == 5:
	""" 3 lvl AG """
	k1_us = 123.9
	k2_us = 2.5
	d_K = 69.3
	capK_us = 0.0067
	dE_K = 1229.0
	#n1 = 1.0
	n1 = 9
	#kag1_us = 354.6
	#kag2_us = 1.21
	#dag_K = 103.3
	kag1_us = 356.0
	#kag2_us = 0.1
	kag2_us = 9.58
	dag_K = 66.4
	capKag_us = 0.000019
	deltaEag_K = 3272.3
	#n3 = 0.1318
	n3 = 1.2
	
elif modelAgLvl == 3:
	""" 2 lvl AG """
	k1_us = 123.9
	k2_us = 2.5
	d_K = 69.3
	capK_us = 0.0067
	dE_K = 1229.0
	#n1 = 1.0
	#n1 = 7
	#kag1_us = 354.6
	#kag2_us = 1.21
	#dag_K = 103.3
	kag1_us = 356.0
	#kag2_us = 0.1
	kag2_us = 0.1
	dag_K = 150
	#n3 = 0.1318
	#n3 = 0.5
	
	#n1 = 6
	#n3 = 0.1
	
	n1 = 9
	n3 = 1.2

elif modelAgLvl == 1:
	""" 1 lvl AG """
	k1_us = 123.9
	k2_us = 2.5
	d_K = 69.3
	capK_us = 0.0067
	dE_K = 1229.0
	#n1 = 1.0
	n1 = 9
	#kag1_us = 354.6
	#kag2_us = 1.21
	#dag_K = 103.3
	kag1_us = 356.0
	#kag2_us = 0.1
	#kag2_us = 9.58
	#dag_K = 66.4
	#capKag_us = 0.000019
	#deltaEag_K = 3272.3
	#n3 = 0.1318
	n3 = 1.2
	

minT = 1.0
maxT = 1000.0
temper = ScintUtil.logFill(minT,maxT)
model =0

if modelAgLvl == 1:
	model = ThreeLvlNonRadTrap(k1_us, k2_us, d_K, capK_us, dE_K, n1, kag1_us, n3)
elif modelAgLvl == 3:
	model = ThreeLvlNonRadTrap2Terms(k1_us, k2_us, d_K, capK_us, dE_K, n1, kag1_us, kag2_us, dag_K, n3)
elif modelAgLvl == 5:
	model = ThreeLvlNonRadTrap3Terms(k1_us, k2_us, d_K, capK_us, dE_K, n1, kag1_us, kag2_us, dag_K, capKag_us, deltaEag_K, n3)

tau_us = []
tauAg_us = []
lyMain = []
lyAg = []
lyTot = []

for i in range(len(temper)):
	param = CalcParameters(model, temper[i])
	tau_us.append(param.tau_us)
	tauAg_us.append(param.tauAg_us)
	lyMain.append(param.ly)
	lyAg.append(param.lyAg)
	lyTot.append(param.lyTot)



"""
###############################################################
# Real Data
###############################################################
"""
treeL = AidaFileReader(aidaPathL, aidaFileL).tree
treeQ = AidaFileReader(aidaPathQ, aidaFileQ).tree

lyDPS = GrabIObject(treeL, "/Light_Yield/LY Corr").dps
aUtil = AidaUtil()
aUtil.sortUp(lyDPS)

mainDecay = ScintUtil.fusionDPS(GrabIObject(treeL,"/Kinetic/Decay_Comp#1").dps, GrabIObject(treeQ,"/Kinetic/Decay_Q_Comp#2").dps)
agDecay = ScintUtil.fusionDPS(GrabIObject(treeL,"/Kinetic/Decay_Comp#2").dps, GrabIObject(treeQ,"/Kinetic/Decay_Q_Comp#3").dps)


"""
###############################################################
# Functions
###############################################################
"""
""" LY Function """
lyFunc=0
if modelAgLvl == 1:
	lyFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/((1.0/kag1_us))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,A,B", "LY Fit 3 Lvl AG")
	lyFunc.setParameter("kag1_us", kag1_us)
	

elif modelAgLvl == 3:
	lyFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/A/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0]))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,kag2_us,dEag_K,A,B", "LY Fit 3 Lvl AG")
	lyFunc.setParameter("kag1_us", kag1_us)
	lyFunc.setParameter("kag2_us", kag2_us)
	lyFunc.setParameter("dEag_K", dag_K)

elif modelAgLvl == 5:
	lyFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("LY Fit 3 Lvl AG", 1, "A*((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0]))*((1.0/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1.0+exp(-1.0*D_K/x[0]))*(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])))-(B/A/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0])+(1.0+exp(-1.0*dEag_K/x[0]))*(1.0/Kag_us)*exp(-1.0*DeltaEag_K/x[0]))))", "k1_us,k2_us,D_K,K_us,DeltaE_K,kag1_us,kag2_us,dEag_K,Kag_us,DeltaEag_K,A,B", "LY Fit 3 Lvl AG")
	lyFunc.setParameter("kag1_us", kag1_us)
	lyFunc.setParameter("kag2_us", kag2_us)
	lyFunc.setParameter("dEag_K", dag_K)
	lyFunc.setParameter("Kag_us", capKag_us)
	lyFunc.setParameter("DeltaEag_K", deltaEag_K)

lyFunc.setParameter("k1_us", k1_us)
lyFunc.setParameter("k2_us", k2_us)
lyFunc.setParameter("D_K", d_K)
lyFunc.setParameter("K_us", capK_us)
lyFunc.setParameter("DeltaE_K", dE_K)
lyFunc.setParameter("A", n1)
lyFunc.setParameter("B", n3)

""" Main Decay Function """
mainFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Kinetic Fit", 1, "(1.0+exp(-1.0*D_K/x[0]))/((1.0/k1_us)+(1.0/k2_us)*exp(-1.0*D_K/x[0])+(1.0/K_us)*exp(-1.0*DeltaE_K/x[0])+(1.0/K_us)*exp(-1.0*(D_K+DeltaE_K)/x[0]))", "k1_us,k2_us,D_K,K_us,DeltaE_K", "Kinetic Fit")
mainFunc.setParameter("k1_us", k1_us)
mainFunc.setParameter("k2_us", k2_us)
mainFunc.setParameter("D_K", d_K)
mainFunc.setParameter("K_us", capK_us)
mainFunc.setParameter("DeltaE_K", dE_K)


""" Afterglow Function """
agFunc = 0
if modelAgLvl == 1:
	agFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Afterglow Fit", 1, "kag1_us", "kag1_us", "Afterglow Fit")
	agFunc.setParameter("kag1_us", kag1_us)
	
elif modelAgLvl == 3:
	agFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Afterglow Fit", 1, "1.0/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0]))", "kag1_us,kag2_us,dEag_K", "Afterglow Fit")
	agFunc.setParameter("kag1_us", kag1_us)
	agFunc.setParameter("kag2_us", kag2_us)
	agFunc.setParameter("dEag_K", dag_K)
		
elif modelAgLvl == 5:
	agFunc = af.createFunctionFactory(af.createTreeFactory().create()).createFunctionFromScript("Afterglow Fit", 1, "1.0/((1.0/kag1_us)+(1.0/kag2_us)*exp(-1.0*dEag_K/x[0])+(1.0+exp(-1.0*dEag_K/x[0]))*(1.0/capKag_us)*exp(-1.0*deltaEag_K/x[0]))", "kag1_us,kag2_us,dEag_K,capKag_us,deltaEag_K", "Afterglow Fit")
	agFunc.setParameter("kag1_us", kag1_us)
	agFunc.setParameter("kag2_us", kag2_us)
	agFunc.setParameter("dEag_K", dag_K)
	agFunc.setParameter("capKag_us", capKag_us)
	agFunc.setParameter("deltaEag_K", deltaEag_K)


"""
###############################################################
# Plotting Part
###############################################################
"""
""" Create the plottable JHep objects from AIDA IObjects"""
tau =P1D(mainDecay)
tau.title = "Main"
tau.symbol = 4
tau.symbolSize = 4
tau.color = colors[0]
tau.errAll = False

tauAg =P1D(agDecay)
tauAg.title = "Afterglow"
tauAg.symbol = 4
tauAg.symbolSize = 4
tauAg.color = colors[1]
tauAg.errAll = False

tauF = F1D(mainFunc, 0.1, 1000)
tauF.title = "Main Fit"
tauF.points = 10000
tauF.color = colors[0]
tauF.penWidth = 2

tauAgF = F1D(agFunc, 0.1, 1000)
tauAgF.title = "Afterglow Fit"
tauAgF.points = 10000
tauAgF.color = colors[1]
tauAgF.penWidth = 2

ly =P1D("LY", temper, lyMain)
ly.symbol = 4
ly.symbolSize = 4
ly.color = colors[0]
ly.errAll = False

lyAg =P1D("LYag", temper, lyAg)
lyAg.symbol = 4
lyAg.symbolSize = 4
lyAg.color = colors[1]
lyAg.errAll = False

lyT =P1D("LYtot", temper, lyTot)
lyT.symbol = 4
lyT.symbolSize = 4
lyT.color = colors[2]
lyT.errAll = False

lyData =P1D(lyDPS)
lyData.title = "Afterglow Fit"
lyData.symbol = 2
lyData.symbolSize = 12
lyData.color = Color.magenta
lyData.errAll = False


""" Create the Keys for the Legend """
lyKey = JaKey(ly.symbol, "Main", 0.20, 0.70, "NDC")
lyKey.color = ly.color
lyKey.keyColor = ly.color
lyKey.penWidth = 2
lyKey.font = legendFont

lyAgKey = JaKey(lyAg.symbol, "Afterglow", 0.20, 0.75, "NDC")
lyAgKey.color = lyAg.color
lyAgKey.keyColor = lyAg.color
lyAgKey.penWidth = 2
lyAgKey.font = legendFont

lyTKey = JaKey(lyT.symbol, "LY Total", 0.20, 0.80, "NDC")
lyTKey.color = lyT.color
lyTKey.keyColor = lyT.color
lyTKey.penWidth = 2
lyTKey.font = legendFont

lyDataKey = JaKey(lyData.symbol, "LY Exp", 0.20, 0.85, "NDC")
lyDataKey.color = lyData.color
lyDataKey.keyColor = lyData.color
lyDataKey.penWidth = 2
lyDataKey.font = legendFont


""" Create the Plotter for the tau"""
tauPlotter = HPlotJa("tau", 600, 600, 1, 1, True)
tauPlotter.setGTitle("Scint Model tau", titleFont)
tauPlotter.antiAlias = False
if showFigures:
	tauPlotter.visible(True)
else:
	tauPlotter.visible(False)

tauPlotter.setPad(1, 1, 0.125, 0.125, 0.75, 0.75)
tauPlotter.setRange(1, 1000, 0.1, 1000)
tauPlotter.setLogScale(0, True)
tauPlotter.setLogScale(1, True)
tauPlotter.ticFont = ticFont
tauPlotter.showStatBox = False
tauPlotter.showKey(False)
tauPlotter.nameX = "Temperature (K)"
tauPlotter.nameY = "Decay Time (&micro;s)"

tauPlotter.draw(tau)
tauPlotter.draw(tauAg)
tauPlotter.draw(tauF)
tauPlotter.draw(tauAgF)
tauPlotter.draw(lyKey)
tauPlotter.draw(lyAgKey)


	
#tauPlotter.draw(shKeyQ)
#tauPlotter.draw(mnKeyQ)
#tauPlotter.draw(lgKeyQ)
#tauPlotter.draw(txt)
	
""" Modifies the JaAxes and the axes labels"""
labX_tau = tauPlotter.getLabelX()[0][0]
labX_tau.font = axisLabelFont
HPlotJaUtil.setLabelLocationAtCenter(0, labX_tau, tauPlotter)
labY_tau = tauPlotter.getLabelY()[0][0]
labY_tau.font = axisLabelFont
HPlotJaUtil.setLabelLocationAtCenter(1, labY_tau, tauPlotter, -0.01)

theAxes_tau = tauPlotter.getPad()
theAxes_tau.setLabelSpace(0, 0.0)
theAxes_tau.setLabelSpace(1, -0.025)
HPlotJaUtil.setTickLabelOneToPower(0, theAxes_tau)
HPlotJaUtil.setTickLabelOneToPower(1, theAxes_tau)

tauPlotter.update()

#""" Saves or not the plot"""
#if writeFigures:
#	for i in range(len(figType)):
#		outFileName = "tauModel1lvlAg" + figType[i]
#		tauPlotter.export(outPath + File.separator + outFileName)
#		print("Writting File : " + outPath + File.separator + outFileName)
#		
#tauPlotter.update()



""" Create the Plotter for the LY"""
lyPlotter = HPlotJa("LY", 600, 600, 1, 1, True)
lyPlotter.setGTitle("Scint Model LY", titleFont)
lyPlotter.antiAlias = False
if showFigures:
	lyPlotter.visible(True)
else:
	lyPlotter.visible(False)

lyPlotter.setPad(1, 1, 0.125, 0.125, 0.75, 0.75)
lyPlotter.setRange(1, 1000, 0.0, 10)
lyPlotter.setLogScale(0, True)
lyPlotter.setLogScale(1, False)
lyPlotter.ticFont = ticFont
lyPlotter.showStatBox = False
lyPlotter.showKey(False)
lyPlotter.nameX = "Temperature (K)"
lyPlotter.nameY = "LY"

lyPlotter.draw(ly)
lyPlotter.draw(lyAg)
lyPlotter.draw(lyT)
lyPlotter.draw(lyData)
lyPlotter.draw(lyKey)
lyPlotter.draw(lyAgKey)
lyPlotter.draw(lyTKey)
lyPlotter.draw(lyDataKey)
	
#lyPlotter.draw(shKeyQ)
#lyPlotter.draw(mnKeyQ)
#lyPlotter.draw(lgKeyQ)
#lyPlotter.draw(txt)
	
""" Modifies the JaAxes and the axes labels"""
labX_ly = lyPlotter.getLabelX()[0][0]
labX_ly.font = axisLabelFont
HPlotJaUtil.setLabelLocationAtCenter(0, labX_ly, lyPlotter)
labY_ly = lyPlotter.getLabelY()[0][0]
labY_ly.font = axisLabelFont
HPlotJaUtil.setLabelLocationAtCenter(1, labY_ly, lyPlotter, -0.01)

theAxes_ly = lyPlotter.getPad()
theAxes_ly.setLabelSpace(0, 0.0)
#theAxes_ly.setLabelSpace(1, -0.025)
theAxes_ly.setLabelSpace(1, 0.0)
HPlotJaUtil.setTickLabelOneToPower(0, theAxes_ly)
HPlotJaUtil.setTickLabelOneToPower(1, theAxes_ly)

lyPlotter.update()

""" Saves or not the plot"""
if writeFigures:
	outFileNameLY = 0
	outFileName = 0
	for i in range(len(figType)):
		if modelAgLvl == 5:
			outFileNameLY = "lyModel3lvlAg" + figType[i]
			outFileName = "tauModel3lvlAg" + figType[i]
		elif modelAgLvl == 3:
			outFileNameLY = "lyModel2lvlAg" + figType[i]
			outFileName = "tauModel2lvlAg" + figType[i]
		elif modelAgLvl == 1:
			outFileNameLY = "lyModel1lvlAg" + figType[i]
			outFileName = "tauModel1lvlAg" + figType[i]
		
		tauPlotter.export(outPath + File.separator + outFileNameLY)
		print("Writting File : " + outPath + File.separator + outFileNameLY)
		lyPlotter.export(outPath + File.separator + outFileName)
		print("Writting File : " + outPath + File.separator + outFileName)

tauPlotter.update()
lyPlotter.update()

print("That's all folks !")

